// src/ir/initializer.rs
extern crate alloc;

use alloc::vec::Vec;

use crate::{
    ir::tensor::{DType, Tensor},
    onnx::TensorProto,
};

#[derive(Debug)]
pub enum InitError {
    UnsupportedDataType(i32),
    TypedFieldUnsupported(DType),
    MissingData,
    SizeMismatch { expect: usize, got: usize },
    Overflow,
}

/// Fallible conversion, returns None on error.
pub fn from_tensor_proto(tp: &TensorProto) -> Option<Tensor> {
    try_from_tensor_proto(tp).ok()
}

/// Convert ONNX TensorProto initializer to internal Tensor.
/// Supports both `raw_data` (preferred) and typed fields fallback.
pub fn try_from_tensor_proto(tp: &TensorProto) -> Result<Tensor, InitError> {
    // resolve dtype
    let dt_i32 = tp.data_type.ok_or(InitError::UnsupportedDataType(-1))?;
    let dtype = onnx_dtype_to_dtype(dt_i32).ok_or(InitError::UnsupportedDataType(dt_i32))?;

    // compute shape and element count
    let shape: Vec<usize> = tp.dims.iter().map(|&d| d as usize).collect();
    let elem_count: usize = if shape.is_empty() {
        1
    } else {
        shape.iter().product()
    };

    // determine required byte size
    let need_bytes = elem_count
        .checked_mul(dtype.size_of())
        .ok_or(InitError::Overflow)?;

    // prefer raw_data when non-empty (little-endian per ONNX)
    if let Some(raw) = tp.raw_data.as_ref() {
        if !raw.is_empty() {
            if raw.len() != need_bytes {
                return Err(InitError::SizeMismatch {
                    expect: need_bytes,
                    got: raw.len(),
                });
            }
            return Ok(Tensor {
                dtype,
                shape,
                data: raw.clone(),
            });
        }
    }

    // fallback to typed fields with strict element-count checks
    let mut buf = alloc::vec![0u8; need_bytes];

    match dtype {
        DType::F32 => {
            fill_buf_from_slice(&mut buf, &tp.float_data, elem_count, 4, |v: f32, out| {
                out.copy_from_slice(&v.to_le_bytes());
            })?;
        }
        DType::F64 => {
            fill_buf_from_slice(&mut buf, &tp.double_data, elem_count, 8, |v: f64, out| {
                out.copy_from_slice(&v.to_le_bytes());
            })?;
        }
        DType::I32 => {
            fill_buf_from_slice(&mut buf, &tp.int32_data, elem_count, 4, |v: i32, out| {
                out.copy_from_slice(&v.to_le_bytes());
            })?;
        }
        DType::I64 => {
            fill_buf_from_slice(&mut buf, &tp.int64_data, elem_count, 8, |v: i64, out| {
                out.copy_from_slice(&v.to_le_bytes());
            })?;
        }
        DType::U64 => {
            // ONNX provides uint64_data for U64; otherwise it should be in raw_data.
            if tp.uint64_data.is_empty() {
                return Err(InitError::MissingData);
            }
            fill_buf_from_slice(&mut buf, &tp.uint64_data, elem_count, 8, |v: u64, out| {
                out.copy_from_slice(&v.to_le_bytes());
            })?;
        }
        DType::Bool => {
            // commonly encoded via int32_data (0/1) or raw_data (1 byte per element).
            if tp.int32_data.len() != elem_count {
                return Err(InitError::SizeMismatch {
                    expect: elem_count,
                    got: tp.int32_data.len(),
                });
            }
            // our internal convention: 1 byte per boolean (0 or 1).
            for (i, &v) in tp.int32_data.iter().enumerate() {
                buf[i] = if v == 0 { 0 } else { 1 };
            }
        }

        // these are typically stored only in raw_data in ONNX initializers.
        // if we got here, raw_data was empty → we currently do not support typed fields for them.
        DType::U32 | DType::U16 | DType::U8 | DType::I16 | DType::I8 | DType::F16 | DType::BF16 => {
            return Err(InitError::TypedFieldUnsupported(dtype));
        }
    }

    Ok(Tensor {
        dtype,
        shape,
        data: buf,
    })
}

/// Generic helper to fill `dst` with LE bytes converted from `src`.
/// Avoids const-generic array sizes in signatures to keep it compiling on stable Rust.
fn fill_buf_from_slice<T: Copy, F>(
    dst: &mut [u8],
    src: &[T],
    need_elems: usize,
    elem_size: usize,
    mut write_le: F,
) -> Result<(), InitError>
where
    // `write_le` must write exactly `elem_size` bytes into `out`.
    F: FnMut(T, &mut [u8]),
{
    if src.len() != need_elems {
        return Err(InitError::SizeMismatch {
            expect: need_elems,
            got: src.len(),
        });
    }
    let need_bytes = need_elems * elem_size;
    if dst.len() != need_bytes {
        return Err(InitError::SizeMismatch {
            expect: need_bytes,
            got: dst.len(),
        });
    }

    let mut off = 0usize;
    for &v in src {
        write_le(v, &mut dst[off..off + elem_size]);
        off += elem_size;
    }
    Ok(())
}

/// Map ONNX TensorProto_DataType (int) to internal DType.
fn onnx_dtype_to_dtype(dt: i32) -> Option<DType> {
    match dt {
        1 => Some(DType::F32),   // FLOAT
        2 => Some(DType::U8),    // UINT8
        3 => Some(DType::I8),    // INT8
        4 => Some(DType::U16),   // UINT16
        5 => Some(DType::I16),   // INT16
        6 => Some(DType::I32),   // INT32
        7 => Some(DType::I64),   // INT64
        8 => None,               // TODO: STRING
        9 => Some(DType::Bool),  // BOOL
        10 => Some(DType::F16),  // FLOAT16
        11 => Some(DType::F64),  // DOUBLE
        12 => Some(DType::U32),  // UINT32
        13 => Some(DType::U64),  // UINT64
        14 | 15 => None,         // TODO: COMPLEX64/128
        16 => Some(DType::BF16), // BFLOAT16
        _ => None,
    }
}

#[allow(dead_code)]
fn dtype_to_onnx_dtype(dt: DType) -> Option<i32> {
    Some(match dt {
        DType::F32 => 1,
        DType::U8 => 2,
        DType::I8 => 3,
        DType::U16 => 4,
        DType::I16 => 5,
        DType::I32 => 6,
        DType::I64 => 7,
        DType::Bool => 9,
        DType::F16 => 10,
        DType::F64 => 11,
        DType::U32 => 12,
        DType::U64 => 13,
        DType::BF16 => 16,
    })
}
